perf(gdn): make GDN kernels compilation batch-size agnostic (support dynamic batch shapes for vLLM integration) - #3649

Resolves conflicts in flashinfer/gdn_decode.py and
flashinfer/gdn_kernels/gdn_decode_mtp.py between this PR's fp32 4D-pool
feature (output_state_indices, use_pool_indexing, in-place 4D writeback)
and main's flashinfer-ai#3649 (batch-size-agnostic compilation) and flashinfer-ai#3502 (BF16
recovery / per-request K, FLA per-token scatter).

Two integration fixes beyond the mechanical conflict resolution:
- Drop B and pool_size from the inline/warp cache keys: flashinfer-ai#3649 made kernel
  compilation batch-size-agnostic and removed them from the compiled
  kernel stubs, so the resolved cache-key tuples must match (was passing
  20 args to an 18-arg stub).
- Make the h0_source dlpack marking conditional on use_pool_indexing:
  main's mark_compact_shape_dynamic(stride_order=(0,1,2)) assumes a
  contiguous 3D flat pool and fails on the PR's 4D strided pool. Use plain
  from_dlpack for the 4D path (strides keyed via pool_strides_key; the
  pool-dim stride is batch-size-independent, so kernel reuse is safe).

Verified on GB200 (sm_100): 32/32 fp32 MTP pool tests pass; benchmark
shows no regression vs main on the standard contiguous path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>

Integrate latest main (incl. flashinfer-ai#3502/flashinfer-ai#3649 GDN refactors) and resolve
conflicts in the GDN decode kernels and tests. Per-batch q/k/v/a/b/o use
mark_layout_dynamic to accept non-compact packed Q/K/V; pool/intermediate
tensors keep mark_compact_shape_dynamic(mode=0) so launchers can derive
constexpr tile counts.

<!-- .github/pull_request_template.md -->

## 📌 Description

#3649 made the GDN decode kernels batch-size agnostic by marking the
per-batch tensors (q/k/v/a/b/o) with CuTe's
`mark_compact_shape_dynamic`, which requires a compact layout. This was
a functional layout coverage regression, and flagged as being
incompatible with an e2e sglang test.

The fix is to mark the per-batch tensors with `mark_layout_dynamic`
instead of `mark_compact_shape_dynamic`.

To verify this fix, this PR adds packed-QKV coverage for unit tests:
q/k/v built as head-dim slices of a fused buffer (the SGLang layout),
asserting results to the contiguous equivalents. These fail on current
main (reproducing the regression) and pass with the fix.

## 🔍 Related Issues

none

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [x] Tests have been added or updated as needed.
- [x] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Refactor**
* Improved dynamic-shape/layout annotations used for cached and compiled
decoding kernels across pretranspose, nontranspose, and MTP paths,
better aligning with fused/packed QKV layouts and ensuring consistent
handling of cached intermediates.

* **Bug Fixes**
* Fixed decoding behavior for packed, non-contiguous Q/K/V tensor views
to produce consistent results (including bit-identical outputs vs
contiguous tensors).

* **Tests**
* Added regression tests covering pretranspose, nontranspose, and MTP
decode with packed/non-contiguous Q/K/V views, including output/state
equivalence checks.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

<!-- .github/pull_request_template.md -->

## 📌 Description
This is the **fp32 sibling of #3268** — the same 4D-pool support, now
for the **fp32 MTP decode path** (`gated_delta_rule_mtp`).
It does two things:
1. **In-place 4D pool writeback (correctness fix).** Previously a 4D
state pool `[pool_size, HV, V, K]` was reshaped to 3D internally. For a
**non-contiguous** pool (e.g. a strided slice of an oversized backing
buffer) `.reshape()` silently materializes a **copy**, the kernel
updates the copy, and the
caller's pool is left untouched → **state updates lost**. The kernel now
reads/writes the pool **in place via native 4D indexing**
(`use_pool_indexing`), so writes land in the caller's tensor with no
silent copy.
2. **Separate read/write slots (new capability).** New optional
`output_state_indices` arg (shape `[B]`): where to **write** the updated
state, separate from where it is **read** (`initial_state_indices`).
Defaults to `initial_state_indices` (write back to the read slot);
**negative entries skip the
writeback** for that batch slot (matching the read-side `-1` padding
semantics). This is what speculative-decoding / MTP-verify needs: read
prior state from one pool slot, write the verified state to another.
```python
out, state = gated_delta_rule_mtp(
q, k, v,
initial_state=pool, # 4D pool, contiguous OR strided
initial_state_indices=read_idx,
output_state_indices=write_idx, # NEW (optional; defaults to read_idx;
-1 = skip)
...
)
```
The rest of the signature is unchanged; the standard contiguous
single-token/MTP path behaves identically.
### Rebased on current `main`
This branch was merged up to current `main`, which required reconciling
the feature with two landed changes:
- **#3649** (batch-size-agnostic GDN compilation): `B`/`pool_size` are
no longer part of the kernel cache key, so the cache keys here were
updated to match the batch-agnostic kernel stubs.
- **#3502** (BF16 recovery / per-request K, FLA per-token scatter): the
writeback guards now compose with `per_token_pool_scatter`.
The non-contiguous 4D pool is passed through `from_dlpack` without
`mark_compact_shape_dynamic` (which assumes a compact 3D layout);
strides are keyed via `pool_strides_key` and the pool-dim stride is
batch-size-independent, so a single compiled kernel is reused across
batch sizes — compilation stays
batch-size-agnostic.

## 🧪 Tests
  
Added to `tests/gdn/test_decode_delta_rule.py`:
- `test_mtp_fp32_state_pool` — pool read/write across seq lengths and
batch sizes, optional separate output indices, intermediate-state
caching. Params: `batch_size ∈ {1,4,16} × seq_len ∈ {2,4} ×
use_separate_output_indices ∈ {F,T} × cache_intermediate_states ∈ {F,T}`
(24 cases).
- `test_mtp_fp32_state_pool_non_contiguous` — strided (non-contiguous)
pools. Params: `batch_size ∈ {1,4} × seq_len ∈ {2,4} × stride_multiplier
∈ {2,3}` (8 cases).




This is the fp32 sibling of PR #3268 — same 4D-pool support, now for the
fp32 MTP decode path.
  
  ### Before vs after — wrapper-side
  
  ```python
  # Before
  out, state = gated_delta_rule_mtp( 
      q, k, v,
      initial_state=pool,            # 4D pool
      initial_state_indices=read_idx,
      ...
  ) 
  # Internally: pool is reshaped to 3D. If pool is non-contiguous,
  # this silently materializes a copy → kernel updates the COPY,
  # original pool is left untouched. Updates lost.
  
  # After
  out, state = gated_delta_rule_mtp( 
      q, k, v,
      initial_state=pool,            # 4D pool, contiguous OR strided
      initial_state_indices=read_idx,
output_state_indices=write_idx, # NEW (optional, defaults to read_idx)
      ...
  )
  # Kernel reads/writes pool in place via native 4D indexing.
  # No silent copy. Writes land in the caller's tensor.
  

Before vs after — memory layout the kernel sees
    
  Contiguous pool (unchanged fast path):
    caller's [pool, HV, V, K]  ─reshape view→  kernel's [pool*HV, V, K]
                                (free, no copy)
    
  Non-contiguous pool (the new path):  
    Before:  caller's strided [pool, HV, V, K]
             ─.reshape() silent copy→  scratch [pool*HV, V, K]
             kernel writes scratch ❌ (caller's  pool unchanged)
             
    After:   caller's strided [pool, HV, V, K]
             ─pass through, use_pool_indexing=True→  kernel
             kernel writes caller's pool in place ✓
  

API change — one new argument
      
  output_state_indices : torch.Tensor, optional (shape [B])
      Where to WRITE the updated state, separate from where you READ.
      Defaults to initial_state_indices (write back to the read slot).
      Negative entries skip the writeback for that batch.
  
  Everything else in the signature behaves identically.
## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Multi-token FP32 decoding: supports pool-backed initial states,
per-batch configurable output-state writeback indices, and correct
handling for non-contiguous pooled layouts. Non-pool single-token decode
behavior remains unchanged.

* **Tests**
* New FP32 tests cover pool read/write across sequence lengths and batch
sizes, optional separate output indices, intermediate-state caching,
non-contiguous pools, and stride variations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Amey Naik <212485788+ameynaik-hub@users.noreply.github.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…ions (#4513)

## 📌 Description

`tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins
many times
over. Two independent reasons, both verified against the cache keys
rather than
assumed:

**1. Batch size is not a compile key.** It has been dynamic since #3649,
so it
reaches the cache only through coarse buckets:

| Path | How `B` enters the key |
|------|------------------------|
| pretranspose (`gdn_decode_pretranspose.py:964`) | not at all |
| bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all —
the tests build a contiguous pool, so `pool_size_key = -1` and
`pool_slot_stride = (-1,)` are B-independent sentinels |
| nontranspose (`gdn_decode_nontranspose.py:725`) | only via
`use_small_batch = B < 32` |
| fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`,
`gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` /
`_get_bf16_mtp_config` |

The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`:
378 of the
file's 817 parametrized cases (46%) but only **42** compile keys,
because `tile_v`
is an explicit monkeypatched axis and the 9 batch sizes contribute
nothing.

The first commit keeps one batch size per bucket. I verified each kept
set
reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a
real
mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel`
silently
dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}).

**2. Intermediate `seq_len` values only re-specialize on `T`.**
`get_mtp_config`
returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for
every
`T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile
config
that T=4 or T=8 does not already cover. `T=2` is kept as the one
structurally
distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2`
branches.

### Effect

| | cases | compile keys (retuned tests) |
|---|---|---|
| before | 817 | 102 |
| after batch collapse | 501 | 102 |
| after `seq_len` trim | 409 | 62 |

Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized).

### Coverage cost

The first commit costs nothing in specialization coverage — the same
cubins still
run, just at fewer runtime batch sizes. The second commit is a
deliberate
reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a
T-specific off-by-one would no longer be caught. It is a separate commit
so it
can be dropped if reviewers would rather keep the full sweep.

## 🔍 Related Issues

Refs #4110 (GDN cold-compile CI time). Complements #4128 and #4444,
which remove
key entries that provably do not reach codegen; this removes test cases
that map
onto keys already covered.

## 🚀 Pull Request Checklist

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

## 🧪 Tests

- [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416
passed in 32m25s**
- [ ] GPU CI for the timing comparison

## Reviewer Notes

- I did not measure a clean before/after wall clock: the ~38 min
baseline I was
working from comes from #4219's description rather than the same
machine, so
I'd rather let CI provide the comparison than quote a number I can't
stand
  behind. The case and key counts above are exact and static.
- Worth noting for #4110 more broadly: cutting 50% of the cases bought
substantially less than 50% of the wall clock, which suggests the
remaining
cost is dominated by compilation and fixed overhead rather than per-case
execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in
#4214)
  as the larger lever.
- The batch sizes kept per test are load-bearing, not arbitrary — each
set is one
representative per config bucket. I left a one-line comment at each site
so
  they don't get "restored" later.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Tests**
* Streamlined decode and MTP test coverage to use representative batch
sizes and sequence lengths.
* Preserved coverage for key thresholds, tile configurations,
transposition modes, precision variants, and sequence-length scenarios.


<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>

…ions (flashinfer-ai#4513)

## 📌 Description

`tests/gdn/test_decode_delta_rule.py` re-ran the same compiled cubins
many times
over. Two independent reasons, both verified against the cache keys
rather than
assumed:

**1. Batch size is not a compile key.** It has been dynamic since flashinfer-ai#3649,
so it
reaches the cache only through coarse buckets:

| Path | How `B` enters the key |
|------|------------------------|
| pretranspose (`gdn_decode_pretranspose.py:964`) | not at all |
| bf16-state wide-vec (`gdn_decode_bf16_state.py:3440`) | not at all —
the tests build a contiguous pool, so `pool_size_key = -1` and
`pool_slot_stride = (-1,)` are B-independent sentinels |
| nontranspose (`gdn_decode_nontranspose.py:725`) | only via
`use_small_batch = B < 32` |
| fp32 / bf16 MTP (`gdn_decode_mtp.py:2500`,
`gdn_decode_bf16_state.py:3813`) | only via `get_mtp_config` /
`_get_bf16_mtp_config` |

The clearest case was `test_gdn_decode_bf16_state_wide_vec_mtp_kernel`:
378 of the
file's 817 parametrized cases (46%) but only **42** compile keys,
because `tile_v`
is an explicit monkeypatched axis and the 9 batch sizes contribute
nothing.

The first commit keeps one batch size per bucket. I verified each kept
set
reproduces the *full* key set, at `NUM_SMS` 108/132/148 — this caught a
real
mistake, where `[1,8,16,32]` for `test_gdn_decode_bf16_state_t1_kernel`
silently
dropped the `HV=64, tile_v=32` key (that test sweeps `HV` ∈ {32,64}).

**2. Intermediate `seq_len` values only re-specialize on `T`.**
`get_mtp_config`
returns an identical `(tile_v, vec_size, ilp_rows, use_smem_v)` set for
every
`T >= 3`, so T=3/5/6/7 compile fresh MTP cubins without covering a tile
config
that T=4 or T=8 does not already cover. `T=2` is kept as the one
structurally
distinct case — it alone reaches the `ilp=8` and `tile_v=16 / ilp=2`
branches.

### Effect

| | cases | compile keys (retuned tests) |
|---|---|---|
| before | 817 | 102 |
| after batch collapse | 501 | 102 |
| after `seq_len` trim | 409 | 62 |

Collected tests go 838 → 416 (409 parametrized + 7 non-parametrized).

### Coverage cost

The first commit costs nothing in specialization coverage — the same
cubins still
run, just at fewer runtime batch sizes. The second commit is a
deliberate
reduction: T=3/5/6/7 still exercise distinct unrolled loop counts, so a
T-specific off-by-one would no longer be caught. It is a separate commit
so it
can be dropped if reviewers would rather keep the full sweep.

## 🔍 Related Issues

Refs flashinfer-ai#4110 (GDN cold-compile CI time). Complements flashinfer-ai#4128 and flashinfer-ai#4444,
which remove
key entries that provably do not reach codegen; this removes test cases
that map
onto keys already covered.

## 🚀 Pull Request Checklist

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

## 🧪 Tests

- [x] `pytest tests/gdn/test_decode_delta_rule.py -q` on H100: **416
passed in 32m25s**
- [ ] GPU CI for the timing comparison

## Reviewer Notes

- I did not measure a clean before/after wall clock: the ~38 min
baseline I was
working from comes from flashinfer-ai#4219's description rather than the same
machine, so
I'd rather let CI provide the comparison than quote a number I can't
stand
  behind. The case and key counts above are exact and static.
- Worth noting for flashinfer-ai#4110 more broadly: cutting 50% of the cases bought
substantially less than 50% of the wall clock, which suggests the
remaining
cost is dominated by compilation and fixed overhead rather than per-case
execution. That points at persistent/AOT CuTe-DSL artifacts (GDN-P1 in
flashinfer-ai#4214)
  as the larger lever.
- The batch sizes kept per test are load-bearing, not arbitrary — each
set is one
representative per config bucket. I left a one-line comment at each site
so
  they don't get "restored" later.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Tests**
* Streamlined decode and MTP test coverage to use representative batch
sizes and sequence lengths.
* Preserved coverage for key thresholds, tile configurations,
transposition modes, precision variants, and sequence-length scenarios.


<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: kahyunnam <kahyunnam@users.noreply.github.com>